#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PDF文字识别工具
递归遍历指定文件夹，提取所有PDF文件中的文字内容并保存到txt文件
"""

import os
import sys
from pathlib import Path
from datetime import datetime

try:
    import pdfplumber
    PDF_LIBRARY = 'pdfplumber'
except ImportError:
    try:
        import PyPDF2
        PDF_LIBRARY = 'PyPDF2'
    except ImportError:
        print("错误：未安装PDF处理库！请运行: pip install pdfplumber")
        sys.exit(1)


def extract_text_with_pdfplumber(pdf_path):
    """使用pdfplumber提取PDF文字"""
    text_content = []
    try:
        with pdfplumber.open(pdf_path) as pdf:
            for page_num, page in enumerate(pdf.pages, 1):
                text = page.extract_text()
                if text:
                    text_content.append(f"\n{'='*60}\n")
                    text_content.append(f"第 {page_num} 页\n")
                    text_content.append(f"{'='*60}\n")
                    text_content.append(text)
        return ''.join(text_content)
    except Exception as e:
        return f"错误：无法读取PDF文件 - {str(e)}"


def extract_text_with_pypdf2(pdf_path):
    """使用PyPDF2提取PDF文字"""
    text_content = []
    try:
        with open(pdf_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            total_pages = len(pdf_reader.pages)
            
            for page_num in range(total_pages):
                page = pdf_reader.pages[page_num]
                text = page.extract_text()
                if text:
                    text_content.append(f"\n{'='*60}\n")
                    text_content.append(f"第 {page_num + 1} 页\n")
                    text_content.append(f"{'='*60}\n")
                    text_content.append(text)
        return ''.join(text_content)
    except Exception as e:
        return f"错误：无法读取PDF文件 - {str(e)}"


def extract_text_from_pdf(pdf_path):
    """从PDF文件中提取文字"""
    if PDF_LIBRARY == 'pdfplumber':
        return extract_text_with_pdfplumber(pdf_path)
    else:
        return extract_text_with_pypdf2(pdf_path)


def find_all_pdf_files(root_folder):
    """递归查找文件夹中的所有PDF文件"""
    pdf_files = []
    root_path = Path(root_folder)
    
    if not root_path.exists():
        print(f"错误：文件夹不存在 - {root_folder}")
        return pdf_files
    
    if not root_path.is_dir():
        print(f"错误：路径不是文件夹 - {root_folder}")
        return pdf_files
    
    # 递归查找所有PDF文件
    for pdf_file in root_path.rglob("*.pdf"):
        if pdf_file.is_file():
            pdf_files.append(pdf_file)
    
    # 同时查找PDF大写扩展名
    for pdf_file in root_path.rglob("*.PDF"):
        if pdf_file.is_file() and pdf_file not in pdf_files:
            pdf_files.append(pdf_file)
    
    return sorted(pdf_files)


def process_pdfs(folder_path, output_file=None):
    """处理文件夹中的所有PDF文件"""
    print(f"开始扫描文件夹: {folder_path}")
    print(f"使用的PDF库: {PDF_LIBRARY}")
    print("-" * 60)
    
    # 查找所有PDF文件
    pdf_files = find_all_pdf_files(folder_path)
    
    if not pdf_files:
        print("未找到任何PDF文件！")
        return
    
    print(f"找到 {len(pdf_files)} 个PDF文件\n")
    
    # 如果没有指定输出文件，使用默认名称
    if output_file is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_file = f"pdf_extracted_text_{timestamp}.txt"
    
    # 提取文字并写入文件
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(f"PDF文字提取结果\n")
        f.write(f"提取时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"源文件夹: {folder_path}\n")
        f.write(f"PDF文件数量: {len(pdf_files)}\n")
        f.write(f"\n{'#'*60}\n\n")
        
        for idx, pdf_path in enumerate(pdf_files, 1):
            print(f"[{idx}/{len(pdf_files)}] 正在处理: {pdf_path.name}")
            
            # 写入文件信息
            f.write(f"\n{'#'*60}\n")
            f.write(f"文件 {idx}: {pdf_path.name}\n")
            f.write(f"完整路径: {pdf_path}\n")
            f.write(f"{'#'*60}\n")
            
            # 提取文字
            text = extract_text_from_pdf(str(pdf_path))
            
            if text.strip():
                f.write(text)
                f.write("\n\n")
            else:
                f.write("（该PDF文件未提取到文字内容）\n\n")
    
    print("-" * 60)
    print(f"完成！所有文字已保存到: {output_file}")
    print(f"输出文件完整路径: {os.path.abspath(output_file)}")


def main():
    """主函数"""
    print("="*60)
    print("PDF文字识别工具")
    print("="*60)
    print()
    
    # 获取文件夹路径
    if len(sys.argv) > 1:
        folder_path = sys.argv[1]
    else:
        folder_path = input("请输入要扫描的文件夹路径: ").strip()
    
    # 去除引号（如果用户输入时带了引号）
    folder_path = folder_path.strip('"').strip("'")
    
    # 获取输出文件名（可选）
    output_file = None
    if len(sys.argv) > 2:
        output_file = sys.argv[2]
    else:
        user_input = input("请输入输出文件名（直接回车使用默认名称）: ").strip()
        if user_input:
            output_file = user_input
    
    print()
    
    # 处理PDF文件
    try:
        process_pdfs(folder_path, output_file)
    except KeyboardInterrupt:
        print("\n\n程序被用户中断")
    except Exception as e:
        print(f"\n发生错误: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()

